package main

import (
	"encoding/json"
	"fmt"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Jeffail/tunny"
	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	dto "github.com/prometheus/client_model/go"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

// Config 服务器配置
type Config struct {
	Port              string        `yaml:"port"`
	HeartbeatInterval time.Duration `yaml:"heartbeat_interval"`
	ConnectionTimeout time.Duration `yaml:"connection_timeout"`
	MaxConnections    int           `yaml:"max_connections"`
	ChannelBufferSize int           `yaml:"channel_buffer_size"`
	CleanupInterval   time.Duration `yaml:"cleanup_interval"`
	InactiveThreshold time.Duration `yaml:"inactive_threshold"`
	BroadcastPoolSize int           `yaml:"broadcast_pool_size"`
}

// Client 代表一个客户端连接
type Client struct {
	ID         string
	DeviceID   string
	Channel    chan string `json:"-"`
	closed     int32       // 原子操作标记
	LastActive int64       // 最后活动时间戳(原子操作)
}

func (c *Client) Close() {
	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		close(c.Channel)
	}
}

func (c *Client) IsClosed() bool {
	return atomic.LoadInt32(&c.closed) == 1
}

func (c *Client) UpdateActivity() {
	atomic.StoreInt64(&c.LastActive, time.Now().Unix())
}

// ClientManager 管理所有客户端连接
type ClientManager struct {
	Clients sync.Map // key: userID, value: *sync.Map(deviceID->*Client)
}

// NewClientManager 创建新的客户端管理器
func NewClientManager() *ClientManager {
	return &ClientManager{}
}

// AddClient 添加客户端
func (cm *ClientManager) AddClient(client *Client) {
	devices, _ := cm.Clients.LoadOrStore(client.ID, &sync.Map{})
	devices.(*sync.Map).Store(client.DeviceID, client)
	client.UpdateActivity()
}

// RemoveClient 移除客户端
func (cm *ClientManager) RemoveClient(client *Client) {
	if devices, ok := cm.Clients.Load(client.ID); ok {
		// 只在确实删除时才关闭连接
		if _, loaded := devices.(*sync.Map).LoadAndDelete(client.DeviceID); loaded {
			client.Close()

			// 检查是否还有其他设备
			var hasOtherDevices bool
			devices.(*sync.Map).Range(func(_, _ interface{}) bool {
				hasOtherDevices = true
				return false
			})
			if !hasOtherDevices {
				cm.Clients.Delete(client.ID)
			}
		}
	}
}

// CloseAllDevicesByID 关闭指定用户的所有设备连接
func (cm *ClientManager) CloseAllDevicesByID(ID string) error {
	if devices, ok := cm.Clients.Load(ID); ok {
		devices.(*sync.Map).Range(func(k, v interface{}) bool {
			client := v.(*Client)
			select {
			case client.Channel <- "close":
				log.Info().Str("userID", ID).Str("deviceID", k.(string)).Msg("Sent close signal")
			case <-time.After(100 * time.Millisecond):
				log.Info().Str("userID", ID).Str("deviceID", k.(string)).Msg("Timeout sending close signal")
			}

			go func() {
				time.Sleep(1 * time.Second)
				if _, loaded := devices.(*sync.Map).LoadAndDelete(k); loaded {
					client.Close()
					if getGaugeValue(ID) > 0 {
						connectionsGauge.WithLabelValues(ID).Dec()
					}
					logConnectionChange(ID, -1)
				}
			}()
			return true
		})
	}
	return nil
}

// GetAllDevicesByID 获取用户的所有设备ID
func (cm *ClientManager) GetAllDevicesByID(ID string) ([]string, bool) {
	var deviceIDs []string
	if devices, ok := cm.Clients.Load(ID); ok {
		devices.(*sync.Map).Range(func(k, _ interface{}) bool {
			deviceIDs = append(deviceIDs, k.(string))
			return true
		})
		return deviceIDs, len(deviceIDs) > 0
	}
	return nil, false
}

// GetAllConnectedIDs 获取所有已连接的用户ID
func (cm *ClientManager) GetAllConnectedIDs() []string {
	var ids []string
	cm.Clients.Range(func(k, _ interface{}) bool {
		ids = append(ids, k.(string))
		return true
	})
	return ids
}

// ApiResponse API响应结构
type ApiResponse struct {
	Status  bool     `json:"ok"`
	Reason  string   `json:"reason"`
	Devices []string `json:"obj,omitempty"`
}

// BroadcastPayload 广播消息负载
type BroadcastPayload struct {
	ID      string
	Message string
}

var (
	clientManager = NewClientManager()
	cfg           = Config{
		Port:              "9091",
		HeartbeatInterval: 25 * time.Second,
		ConnectionTimeout: 10 * time.Minute,
		MaxConnections:    10000,
		ChannelBufferSize: 100,
		CleanupInterval:   5 * time.Minute,
		InactiveThreshold: 2 * time.Hour,
		BroadcastPoolSize: 100,
	}
	broadcastPool *tunny.Pool

	// Prometheus指标
	connectionsGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Name: "sse_active_connections",
		Help: "Number of active SSE connections",
	}, []string{"user_id"})

	messagesCounter = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "sse_messages_total",
		Help: "Total number of SSE messages sent",
	}, []string{"type"})
)

func init() {
	// 初始化日志
	zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})

	// 重置指标
	connectionsGauge.Reset()
	messagesCounter.Reset()

	// 初始化广播池
	broadcastPool = tunny.NewFunc(cfg.BroadcastPoolSize, func(payload interface{}) interface{} {
		p := payload.(*BroadcastPayload)
		if devices, ok := clientManager.Clients.Load(p.ID); ok {
			devices.(*sync.Map).Range(func(_, v interface{}) bool {
				client := v.(*Client)
				if !client.IsClosed() {
					select {
					case client.Channel <- p.Message:
						messagesCounter.WithLabelValues("broadcast").Inc()
						client.UpdateActivity()
					default:
						log.Info().Str("clientID", client.ID).Msg("Client channel full")
					}
				}
				return true
			})
		}
		return nil
	})

	// 启动清理协程
	go startCleanupRoutine()
}

// 获取当前指标值
func getGaugeValue(userID string) int {
	metric := &dto.Metric{}
	if err := connectionsGauge.WithLabelValues(userID).Write(metric); err != nil {
		return -1
	}
	return int(*metric.Gauge.Value)
}

// 记录连接数变化
func logConnectionChange(userID string, change int) {
	log.Debug().
		Str("userID", userID).
		Int("change", change).
		Int("newCount", getGaugeValue(userID)).
		Msg("Connection count changed")
}

func startCleanupRoutine() {
	ticker := time.NewTicker(cfg.CleanupInterval)
	for range ticker.C {
		clientManager.Clients.Range(func(userID, devices interface{}) bool {
			devices.(*sync.Map).Range(func(deviceID, client interface{}) bool {
				c := client.(*Client)
				lastActive := time.Unix(atomic.LoadInt64(&c.LastActive), 0)
				if time.Since(lastActive) > cfg.InactiveThreshold {
					log.Info().
						Str("userID", c.ID).
						Str("deviceID", c.DeviceID).
						Time("lastActive", lastActive).
						Msg("Cleaning up inactive connection")
					// 确保只删除一次
					if _, loaded := devices.(*sync.Map).LoadAndDelete(deviceID); loaded {
						c.Close()
						if getGaugeValue(c.ID) > 0 {
							connectionsGauge.WithLabelValues(c.ID).Dec()
						}
						logConnectionChange(c.ID, -1)
					}
				}
				return true
			})
			return true
		})
	}
}

func sseHandler(c *gin.Context) {
	clientID := c.Query("id")
	deviceID := c.Query("device_id")
	if clientID == "" {
		c.JSON(400, ApiResponse{
			Status: false,
			Reason: "Client ID is required",
		})
		return
	}

	// 设置SSE头
	c.Writer.Header().Set("Content-Type", "text/event-stream")
	c.Writer.Header().Set("Cache-Control", "no-cache")
	c.Writer.Header().Set("Connection", "keep-alive")

	// 创建新客户端
	client := &Client{
		ID:       clientID,
		DeviceID: deviceID,
		Channel:  make(chan string, cfg.ChannelBufferSize),
	}
	client.UpdateActivity()

	// 注册客户端
	clientManager.AddClient(client)
	connectionsGauge.WithLabelValues(clientID).Inc()
	logConnectionChange(clientID, 1)

	defer func() {
		clientManager.RemoveClient(client)
		if getGaugeValue(clientID) > 0 {
			connectionsGauge.WithLabelValues(clientID).Dec()
		}
		logConnectionChange(clientID, -1)
	}()

	log.Info().
		Str("clientID", clientID).
		Str("deviceID", deviceID).
		Msg("New client connected")

	// 设置心跳和超时
	heartbeat := time.NewTicker(cfg.HeartbeatInterval)
	defer heartbeat.Stop()

	timeout := time.NewTimer(cfg.ConnectionTimeout)
	defer timeout.Stop()

	lastHeartbeat := time.Now() // 记录最后心跳时间

	// 主事件循环
	for {
		select {
		case msg := <-client.Channel:

			log.Info().
				Str("clientID", clientID).
				Str("deviceID", deviceID).
				Msg("send msg:" + msg)

			if msg == "close" {
				if _, err := fmt.Fprintf(c.Writer, "event: close\ndata: Connection closed by server\n\n"); err != nil {
					return
				}
				c.Writer.Flush()
				return
			}

			if strings.TrimSpace(msg) == "" {
				log.Debug().Msg("忽略空消息")
				return
			}

			if _, err := fmt.Fprintf(c.Writer, "data: %s\n\n", msg); err != nil {
				return
			}
			c.Writer.Flush()
			client.UpdateActivity()

		case <-heartbeat.C:
			now := time.Now()
			if now.Sub(lastHeartbeat) <= cfg.HeartbeatInterval*6/5 {
				lastHeartbeat = now
			} else {
				log.Info().
					Str("clientID", clientID).
					Str("deviceID", deviceID).
					Msg("heartbeat time bug")
				continue
			}

			if client.IsClosed() {
				log.Info().
					Str("clientID", clientID).
					Str("deviceID", deviceID).
					Msg("Client already closed, skip heartbeat")
				return
			}

			if _, err := fmt.Fprintf(c.Writer, "event: heartbeat\ndata: %d\n\n", now.Unix()); err != nil {
				log.Info().
					Str("clientID", clientID).
					Str("deviceID", deviceID).
					Msg("Heartbeat failed, client disconnected")
				return
			}

			log.Info().
				Str("clientID", clientID).
				Str("deviceID", deviceID).
				Msg("heartbeat push log" + now.String())
			lastHeartbeat = now
			c.Writer.Flush()
			messagesCounter.WithLabelValues("heartbeat").Inc()

		case <-timeout.C:
			log.Info().
				Str("clientID", clientID).
				Str("deviceID", deviceID).
				Msg("Connection timeout, closing")
			return

		case <-c.Request.Context().Done():
			log.Info().
				Str("clientID", clientID).
				Str("deviceID", deviceID).
				Msg("Client disconnected")
			return
		}

		timeout.Reset(cfg.ConnectionTimeout)
	}
}

func apiHandler(c *gin.Context) {
	var payload struct {
		ID       string          `json:"id"`
		DeviceID string          `json:"device_id,omitempty"`
		Message  json.RawMessage `json:"message"`
	}

	if err := c.ShouldBindJSON(&payload); err != nil || payload.ID == "" || len(payload.Message) == 0 {
		c.JSON(400, ApiResponse{
			Status: false,
			Reason: "Invalid request payload",
		})
		return
	}

	messageStr := string(payload.Message)
	log.Debug().Str("message", messageStr).Msg("Received message")

	var devicesSent []string

	if payload.DeviceID != "" {
		// 定向发送到特定设备
		if devices, ok := clientManager.Clients.Load(payload.ID); ok {
			if client, ok := devices.(*sync.Map).Load(payload.DeviceID); ok {
				client := client.(*Client)
				if !client.IsClosed() {
					select {
					case client.Channel <- messageStr:
						devicesSent = append(devicesSent, payload.DeviceID)
						messagesCounter.WithLabelValues("direct").Inc()
						client.UpdateActivity()
					default:
						log.Info().
							Str("clientID", payload.ID).
							Str("deviceID", payload.DeviceID).
							Msg("Failed to send message to client")
					}
				}
			}
		}
	} else {
		// 广播到所有设备
		broadcastPool.Process(&BroadcastPayload{
			ID:      payload.ID,
			Message: messageStr,
		})
		if devices, ok := clientManager.GetAllDevicesByID(payload.ID); ok {
			devicesSent = devices
		}
	}

	if len(devicesSent) == 0 {
		c.JSON(200, gin.H{
			"ok":     false,
			"reason": "暂无设备在线",
		})
		return
	}

	c.JSON(200, ApiResponse{
		Status:  true,
		Reason:  "Message delivered",
		Devices: devicesSent,
	})
}

func getClientsHandler(c *gin.Context) {
	ID := c.Query("id")
	if ID == "" {
		c.JSON(400, ApiResponse{
			Status: false,
			Reason: "ID is required",
		})
		return
	}

	devices, exists := clientManager.GetAllDevicesByID(ID)
	if !exists {
		if err := clientManager.CloseAllDevicesByID(ID); err != nil {
			log.Error().Err(err).Str("userID", ID).Msg("Failed to close all devices")
		}
		c.JSON(200, gin.H{
			"ok":     false,
			"reason": "暂无设备在线",
		})
		return
	}

	c.JSON(200, gin.H{
		"ok":  true,
		"obj": devices,
	})
}

func closeClientHandler(c *gin.Context) {
	ID := c.Query("id")
	if ID == "" {
		c.JSON(400, ApiResponse{
			Status: false,
			Reason: "ID is required",
		})
		return
	}

	if err := clientManager.CloseAllDevicesByID(ID); err != nil {
		c.JSON(500, ApiResponse{
			Status: false,
			Reason: err.Error(),
		})
		return
	}

	c.JSON(200, ApiResponse{
		Status: true,
		Reason: "Connections closed",
	})
}

func getAllConnectedIDsHandler(c *gin.Context) {
	ids := clientManager.GetAllConnectedIDs()
	c.JSON(200, gin.H{
		"ok":  true,
		"ids": ids,
	})
}

func main() {
	// 设置Gin模式
	gin.SetMode(gin.ReleaseMode)
	r := gin.Default()

	// 添加CORS中间件
	r.Use(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(204)
			return
		}

		c.Next()
	})

	// 注册路由
	r.GET("/events", sseHandler)
	r.POST("/send", apiHandler)
	r.GET("/clients", getClientsHandler)
	r.GET("/closeClient", closeClientHandler)
	r.GET("/connectedIDs", getAllConnectedIDsHandler)
	r.GET("/metrics", gin.WrapH(promhttp.Handler()))

	// 启动服务器
	log.Info().
		Str("port", cfg.Port).
		Msg("Starting SSE server with Gin")

	if err := r.Run(":" + cfg.Port); err != nil {
		log.Fatal().Err(err).Msg("Server failed to start")
	}
}
